import os
import pandas as pd
import matplotlib.pyplot as plt

# 读取数据
script_dir = os.path.dirname(os.path.abspath(__file__))
file_path = os.path.join(script_dir, '..', '..', 'data', 'swimmer_pid_ppo_data.txt')
data = pd.read_csv(file_path, delimiter=',', header=None)

time = data[0] * 0.001  # 转换为秒
actual_angles = data[[1, 2]]
target_angles = data[[3, 4]]
tracking_error = data[[5, 6]]
control_output = data[[13, 14]]
pid_params = data[[7, 8, 9, 10, 11, 12]]

labels = ['Joint1', 'Joint2']
ylabels_angles = ['angle (rad)', 'angle (rad)']
ylabels_error = ['error (rad)', 'error (rad)']
ylabels_control = ['control output (N·m)', 'control output (N·m)']
ylabels_pid = ['Kp1', 'Ki1', 'Kd1', 'Kp2', 'Ki2', 'Kd2']

# 绘制实际角度和目标角度
fig, axs = plt.subplots(2, 1, figsize=(12, 12))
for i in range(2):
    axs[i].plot(time, actual_angles.iloc[:, i], label=f'Actual {labels[i]}')
    axs[i].plot(time, target_angles.iloc[:, i], label=f'Target {labels[i]}')
    axs[i].set_xlabel('time (s)', fontsize=28)
    axs[i].set_ylabel(ylabels_angles[i], fontsize=28)
    axs[i].legend(fontsize=16)
    axs[i].grid()
    axs[i].tick_params(axis='both', which='major', labelsize=28)
plt.tight_layout()
# plt.savefig('Actual_vs_Target_Angles.eps', dpi=600, format='eps')

# 绘制跟踪误差
fig, axs = plt.subplots(2, 1, figsize=(12, 12))
for i in range(2):
    axs[i].plot(time, tracking_error.iloc[:, i], label=f'Tracking Error {labels[i]}')
    axs[i].set_xlabel('time (s)', fontsize=28)
    axs[i].set_ylabel(ylabels_error[i], fontsize=28)
    axs[i].legend(fontsize=16)
    axs[i].grid()
    axs[i].tick_params(axis='both', which='major', labelsize=28)
plt.tight_layout()
# plt.savefig('Tracking_Error.eps', dpi=600, format='eps')

# 绘制控制输出
fig, axs = plt.subplots(2, 1, figsize=(12, 12))
for i in range(2):
    axs[i].plot(time, control_output.iloc[:, i], label=f'Control Output {labels[i]}')
    axs[i].set_xlabel('time (s)', fontsize=28)
    axs[i].set_ylabel(ylabels_control[i], fontsize=28)
    axs[i].legend(fontsize=16)
    axs[i].grid()
    axs[i].tick_params(axis='both', which='major', labelsize=28)
plt.tight_layout()
# plt.savefig('Control_Output.eps', dpi=600, format='eps')

# 绘制 PID 参数
fig, axs = plt.subplots(6, 1, figsize=(12, 18))
for i in range(6):
    axs[i].plot(time, pid_params.iloc[:, i], label=ylabels_pid[i])
    axs[i].set_xlabel('time (s)', fontsize=28)
    axs[i].set_ylabel(ylabels_pid[i], fontsize=28)
    axs[i].legend(fontsize=16)
    axs[i].grid()
    axs[i].tick_params(axis='both', which='major', labelsize=28)
plt.tight_layout()
# plt.savefig('PID_Params.eps', dpi=600, format='eps')

plt.show()
